package com.spring.ai.tutorial.graph.agent.node;

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.AsyncNodeAction;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.alibaba.cloud.ai.graph.agent.ReflectAgent;
import com.alibaba.cloud.ai.graph.node.LlmNode;
import com.alibaba.cloud.ai.graph.state.strategy.AppendStrategy;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.UnaryOperator;

/**
 * @author yingzi
 * @since 2025/8/25
 */

public class JudgeGraphNode implements NodeAction {

    private static final Logger logger = LoggerFactory.getLogger(JudgeGraphNode.class);

    private final LlmNode llmNode;

    private final String NODE_ID = "judge_node";

    private SystemPromptTemplate systemPromptTemplate;

    private static final String CLASSIFIER_PROMPT_TEMPLATE = """
					You are a teacher grading a student's essay submission. Provide detailed feedback and revision suggestions for the essay.

					Your feedback should cover the following aspects:

					- Length : Is the essay sufficiently developed? Does it meet the required length or need expansion/shortening?
					- Depth : Are the ideas well-developed? Is there sufficient analysis, evidence, or explanation?
					- Structure : Is the organization logical and clear? Are the introduction, transitions, and conclusion effective?
					- Style and Tone : Is the writing style appropriate for the purpose and audience? Is the tone consistent and professional?
					- Language Use : Are vocabulary, grammar, and sentence structure accurate and varied?
					- Focus only on providing actionable suggestions for improvement. Do not include grades, scores, or overall summary evaluations.

					Please respond in Chinese .
				""";

    public JudgeGraphNode(ChatClient chatClient) {
        this.systemPromptTemplate = new SystemPromptTemplate(CLASSIFIER_PROMPT_TEMPLATE);
        this.llmNode = LlmNode.builder()
                .chatClient(chatClient)
                .systemPromptTemplate(systemPromptTemplate.render())
                .messagesKey(ReflectAgent.MESSAGES)
                .build();

    }

    public static Builder builder() {
        return new Builder();
    }

    @Override
    public Map<String, Object> apply(OverAllState allState) throws Exception {
        logger.info("judge_node is running.");

        List<Message> messages = (List<Message>) allState.value(ReflectAgent.MESSAGES).get();

        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> keyStrategyHashMap = new HashMap<>();

            keyStrategyHashMap.put(ReflectAgent.MESSAGES, new AppendStrategy());
            return keyStrategyHashMap;
        };

        StateGraph stateGraph = new StateGraph(keyStrategyFactory)
                .addNode(this.NODE_ID, AsyncNodeAction.node_async(llmNode))
                .addEdge(StateGraph.START, this.NODE_ID)
                .addEdge(this.NODE_ID, StateGraph.END);

        CompiledGraph compile = stateGraph.compile();

        OverAllState invokeState = compile.call(Map.of(ReflectAgent.MESSAGES, messages)).get();

        UnaryOperator<List<Message>> convertLastToUserMessage = messageList -> {
            int size = messageList.size();
            if (size == 0)
                return messageList;
            Message last = messageList.get(size - 1);
            messageList.set(size - 1, new UserMessage(last.getText()));
            return messageList;
        };

        List<Message> reactMessages = (List<Message>) invokeState.value(ReflectAgent.MESSAGES).orElseThrow();
        convertLastToUserMessage.apply(reactMessages);

        return Map.of(ReflectAgent.MESSAGES, reactMessages);

    }

    public static class Builder {

        private ChatClient chatClient;

        public JudgeGraphNode.Builder chatClient(ChatClient chatClient) {
            this.chatClient = chatClient;
            return this;
        }

        public JudgeGraphNode build() {
            if (chatClient == null) {
                throw new IllegalArgumentException("ChatClient must be provided");
            }
            return new JudgeGraphNode(chatClient);
        }

    }


}
